import sys
import time
from MyTT import EMA
import numpy as np
import pandas as pd
from tqsdk import TqApi, TqAuth, TqAccount, TqSim, TqKq
from db_manager import connect_db, get_monitored_products
from threading import Thread

# 初始化API
tq_kq = TqKq()
api = TqApi(account=tq_kq, auth=TqAuth("cps168", "alibaba"))
def run_strategy(symbols, lots):
    """主策略函数"""
    running = True
    active_symbols = set()    
    try:        
        # 运行策略主逻辑
        run_all_symbols(api, symbols, lots, running, active_symbols)
        while True:
            api.wait_update(deadline=time.time() + 1)
    except Exception as e:
        print(f"策略运行出错: {e}")
    finally:
        running = False
        # 确保所有线程完成
        while len(active_symbols) > 0:
            time.sleep(0.1)
        # 在主线程关闭API
        if api:
            api.close()

def run_all_symbols(api, symbols, lots, running, active_symbols):
    """处理所有品种"""
    print(f"开始处理品种列表: {symbols}")
    try:
        for symbol in symbols:
            Thread(target=process_product, args=(api, symbol, lots, running, active_symbols)).start()
            
    except Exception as e:
        print(f"策略执行出错: {e}")

def process_product(api, symbol, lots, running, active_symbols):
    """单个品种处理函数"""
    active_symbols.add(symbol)
    try:  
        quote = api.get_quote(symbol)
        klines = api.get_kline_serial(symbol, 1*60, 500)
        
        while running:
            # 等待行情更新
            api.wait_update(deadline=time.time() + 1)
            time.sleep(1)            
            if api.is_changing(quote):
                print(f"{time.strftime('%X')} {symbol} 行情已更新，当前价格: {quote.last_price}")
                q = (3*klines.close + klines.low + klines.open + klines.high) / 6
                ma_values = calculate_ma(q)
                short_ma = ma_values['ma']
                ema_7 = ma_values['ema']
                
                position = api.get_position(symbol)
                
                # 平空仓逻辑
                if short_ma[-2] > ema_7[-2]:  
                    if position.pos_short > 0:
                        api.insert_order(symbol, direction="BUY", offset="CLOSE", 
                                        volume=position.pos_short)
                # 开多仓逻辑
                if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] < ema_7[-3] and short_ma[-2] > ema_7[-2]:
                    api.insert_order(symbol, direction="BUY", offset="OPEN", 
                                    volume=lots)
                    
                # 平多仓逻辑
                if short_ma[-2] < ema_7[-2]:  
                    if position.pos_long > 0:
                        api.insert_order(symbol, direction="SELL", offset="CLOSE", 
                                        volume=position.pos_long)
                # 开空仓逻辑
                if position.pos_long == 0 and position.pos_short == 0 and short_ma[-3] > ema_7[-3] and short_ma[-2] < ema_7[-2]:
                    api.insert_order(symbol, direction="SELL", offset="OPEN", 
                                    volume=lots)

    except Exception as e:
        #print(f"Error processing symbol {symbol}: {e}")  
        pass                      
    finally:
        active_symbols.remove(symbol)

def calculate_ma(q):
    """自定义权重计算"""
    weights = np.arange(26, 0, -1)
    ma = []
    q_array = np.array(q)
    for i in range(len(q_array)):
        if i < 25:
            ma.append(np.nan)
            continue
        window = q_array[i-25:i+1]
        ma.append(np.dot(window, weights) / 351)
    
    ema_7 = EMA(pd.Series(ma), 7)
    return {'ma': ma, 'ema': ema_7.tolist()}